#![allow(dead_code)]
use crate::ops::ops_error::{OpsError, OpsResult};
use alloc::vec;
use alloc::vec::Vec;

/// Compute NumPy-style broadcasted shape for two input shapes (no std/itertools).
pub fn broadcast_shape(left_shape: &[usize], right_shape: &[usize]) -> OpsResult<Vec<usize>> {
    let left_rank = left_shape.len();
    let right_rank = right_shape.len();
    let output_rank = if left_rank > right_rank {
        left_rank
    } else {
        right_rank
    };
    let mut output_shape = vec![0usize; output_rank];

    // align from the tail (right-aligned)
    let mut left_index = left_rank as isize - 1;
    let mut right_index = right_rank as isize - 1;
    let mut out_index = output_rank as isize - 1;

    while out_index >= 0 {
        let left_dim = if left_index >= 0 {
            left_shape[left_index as usize]
        } else {
            1
        };
        let right_dim = if right_index >= 0 {
            right_shape[right_index as usize]
        } else {
            1
        };

        let dim = if left_dim == right_dim {
            left_dim
        } else if left_dim == 1 {
            right_dim
        } else if right_dim == 1 {
            left_dim
        } else {
            return Err(OpsError::ShapeMismatch("incompatible for broadcasting"));
        };

        output_shape[out_index as usize] = dim;

        left_index -= 1;
        right_index -= 1;
        out_index -= 1;
    }

    Ok(output_shape)
}

/// Row-major contiguous strides for a given shape.
pub fn contiguous_strides(shape: &[usize]) -> Vec<isize> {
    let rank = shape.len();
    let mut strides = vec![0isize; rank];
    let mut accumulator: isize = 1;
    for index in 0..rank {
        let rev_index = rank - 1 - index;
        strides[rev_index] = accumulator;
        accumulator *= shape[rev_index] as isize;
    }
    strides
}

/// Compute strides that align `source_shape` to `output_shape`; broadcasted dims get stride = 0.
pub fn broadcast_aligned_strides(
    source_shape: &[usize],
    output_shape: &[usize],
) -> OpsResult<Vec<isize>> {
    let source_rank = source_shape.len();
    let output_rank = output_shape.len();
    let padding = if output_rank > source_rank {
        output_rank - source_rank
    } else {
        0
    };

    // contiguous strides of the source
    let base_strides = contiguous_strides(source_shape);
    let mut aligned_strides = vec![0isize; output_rank];

    // build stride for each dimension from left to right.
    for dim_index in 0..output_rank {
        let source_dim = if dim_index < padding {
            1
        } else {
            source_shape[dim_index - padding]
        };
        let out_dim = output_shape[dim_index];

        if source_dim == out_dim {
            aligned_strides[dim_index] = if dim_index < padding {
                0
            } else {
                base_strides[dim_index - padding]
            };
        } else if source_dim == 1 {
            // broadcasted dimension: fixed cost
            aligned_strides[dim_index] = 0;
        } else {
            return Err(OpsError::ShapeMismatch("broadcast stride alignment failed"));
        }
    }
    Ok(aligned_strides)
}

/// Increment an "odometer"-style coordinate vector; returns false when finished.
pub fn advance_coords(coordinates: &mut [isize], shape: &[usize]) -> bool {
    let rank = coordinates.len();
    if rank == 0 {
        return false;
    }
    let mut k = rank as isize - 1;
    loop {
        coordinates[k as usize] += 1;
        if coordinates[k as usize] < shape[k as usize] as isize {
            return true;
        }
        coordinates[k as usize] = 0;
        if k == 0 {
            return false;
        }
        k -= 1;
    }
}
